package com.bitrate.dify.client.impl;

import cn.hutool.http.Header;
import com.bitrate.dify.client.DifyClient;
import com.bitrate.dify.config.DifyProperties;
import com.bitrate.dify.enums.DifyUrl;
import com.bitrate.dify.enums.ResponseMode;
import com.bitrate.dify.request.ChatCompletionRequest;
import com.bitrate.dify.response.ChatCompletionResponse;
import com.bitrate.dify.util.JsonUtils;
import com.google.gson.Gson;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import okhttp3.internal.sse.RealEventSource;
import okhttp3.sse.EventSourceListener;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.Objects;

@Slf4j
@Service
public class DifyClientImpl implements DifyClient {

    @Autowired
    private DifyProperties difyProperties;

    @Autowired
    private OkHttpClient okHttpClient;

    public static final String BEARER = "Bearer ";

    public static final MediaType APPLICATION_JSON = MediaType.get("application/json; charset=utf-8");

    @Override
    public void streamChatCompletions(ChatCompletionRequest request, EventSourceListener listener) {
        request.setResponse_mode(ResponseMode.STREAMING.getMode());
        Request clientRequest = createRequest(DifyUrl.chat_messages, RequestBody.create(APPLICATION_JSON, JsonUtils.toJson(request)));
        RealEventSource realEventSource = new RealEventSource(clientRequest, listener);
        realEventSource.connect(okHttpClient);
    }

    @Override
    public ChatCompletionResponse blockChatCompletions(ChatCompletionRequest request) {
        request.setResponse_mode(ResponseMode.BLOCKING.getMode());
        log.info("request: {}", JsonUtils.toJson(request));
        Request clientRequest = createRequest(DifyUrl.chat_messages, RequestBody.create(APPLICATION_JSON, JsonUtils.toJson(request)));
        Response response = null;
        try {
            response = okHttpClient.newCall(clientRequest).execute();
            String responseBody = response.body().string();
            ChatCompletionResponse chatCompletionResponse = new Gson().fromJson(responseBody, ChatCompletionResponse.class);
            log.info("responseBody: {}", responseBody);
            return Objects.requireNonNull(chatCompletionResponse);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }


    private Request createRequest(DifyUrl difyUrl, RequestBody requestBody) {
        String requestURI = difyProperties.getDomain().concat(difyUrl.getSuffix());
        String apiKey = difyProperties.getApiKey();
        return new Request.Builder().url(requestURI).header(Header.AUTHORIZATION.name(), BEARER.concat(apiKey)).post(requestBody).build();
    }
}
